{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8c0a98ed-f0cc-437e-bd0b-ba521bcd52f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import models\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8accd3ff-25ec-49fd-aa03-b7468bcdb73a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convrelu(in_channels, out_channels, kernel, padding):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n",
    "        nn.ReLU(inplace=True)\n",
    "    )\n",
    "class ResNetUNet(nn.Module):\n",
    "    def __init__(self, n_class):\n",
    "        super().__init__()\n",
    "\n",
    "        self.base_model = models.resnet18(pretrained=True)\n",
    "        self.base_layers = list(self.base_model.children())\n",
    "\n",
    "        self.layer0 = nn.Sequential(*self.base_layers[:3])\n",
    "        self.layer0_1x1= convrelu(64, 64, 1, 0)\n",
    "        self.layer1 = nn.Sequential(*self.base_layers[3:5])\n",
    "        self.layer1_1x1= convrelu(64, 64, 1, 0)\n",
    "        self.layer2 = self.base_layers[5]\n",
    "        self.layer2_1x1= convrelu(128, 128, 1, 0)\n",
    "        self.layer3 = self.base_layers[6]\n",
    "        self.layer3_1x1= convrelu(256, 256, 1, 0)\n",
    "        self.layer4 = self.base_layers[7]\n",
    "        self.layer4_1x1= convrelu(512, 512, 1, 0)\n",
    "\n",
    "        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "        self.conv_up3 = convrelu(256 + 512, 512, 3, 1)\n",
    "        self.conv_up2 = convrelu(128 + 512, 256, 3, 1)\n",
    "        self.conv_up1 = convrelu(64 + 256, 256, 3, 1)\n",
    "        self.conv_up0 = convrelu(64 + 256, 128, 3, 1)\n",
    "        self.conv_original_size0 = convrelu(3, 64, 3, 1)\n",
    "        self.conv_original_size1 = convrelu(64, 64, 3, 1)\n",
    "        self.conv_original_size2 = convrelu(64 + 128, 64, 3, 1)\n",
    "        self.conv_last = nn.Conv2d(64, n_class, 1)\n",
    "\n",
    "    def forward(self, input):\n",
    "        x_original = self.conv_original_size0(input)\n",
    "        x_original = self.conv_original_size1(x_original)\n",
    "\n",
    "        layer0 = self.layer0(input)\n",
    "        layer1 = self.layer1(layer0)\n",
    "        layer2 = self.layer2(layer1)\n",
    "        layer3 = self.layer3(layer2)\n",
    "        layer4 = self.layer4(layer3)\n",
    "\n",
    "        layer4 = self.layer4_1x1(layer4)\n",
    "        x = self.upsample(layer4)\n",
    "        layer3 = self.layer3_1x1(layer3)\n",
    "        x = torch.cat([x, layer3], dim=1)\n",
    "        x = self.conv_up3(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer2 = self.layer2_1x1(layer2)\n",
    "        x = torch.cat([x, layer2], dim=1)\n",
    "        x = self.conv_up2(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer1 = self.layer1_1x1(layer1)\n",
    "        x = torch.cat([x, layer1], dim=1)\n",
    "        x = self.conv_up1(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer2 = self.layer0_1x1(layer0)\n",
    "        x = torch.cat([x, layer2], dim=1)\n",
    "        x = self.conv_up0(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        x = torch.cat([x, x_original], dim=1)\n",
    "        x = self.conv_original_size2(x)\n",
    "\n",
    "        out = self.conv_last(x)\n",
    "\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "39dcaf9d-9ae8-4b78-9c2d-3fcc05304cd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "#图像预处理\n",
    "transform= transforms.Compose([\n",
    "    transforms.Resize((256, 256)),\n",
    "    transforms.ToTensor()])\n",
    "\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, img_dir, mask_dir, transform=None):\n",
    "        #图像文件夹路径\n",
    "        self.img_dir = img_dir\n",
    "        #掩码文件夹路径\n",
    "        self.mask_dir = mask_dir\n",
    "        #预处理操作\n",
    "        self.transform = transform\n",
    "        #获取图像文件夹内所有.jpg文件\n",
    "        self.image_filenames = [f for f in os.listdir(img_dir) if f.endswith('.jpg')]\n",
    "\n",
    "    def __len__(self):\n",
    "        #返回数据集样本总数\n",
    "        return len(self.image_filenames)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        #获取文件名\n",
    "        img_name = self.image_filenames[idx]\n",
    "        #将图像和掩码的文件名与对应的目录路径组合，形成完整的文件路径\n",
    "        img_path = os.path.join(self.img_dir, img_name)\n",
    "        mask_path = os.path.join(self.mask_dir, img_name)\n",
    "\n",
    "        #加载图像和掩码，转换图像格式（RGB和灰度格式）\n",
    "        image = Image.open(img_path).convert(\"RGB\")\n",
    "        mask = Image.open(mask_path).convert(\"L\")\n",
    "\n",
    "        #应用预处理\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "            mask = self.transform(mask)\n",
    "            #分割掩码二值化，非零值转换为1.0，零值保持0.0；统一标签格式。适配损失函数\n",
    "            mask = (mask > 0).float()\n",
    "        return image, mask\n",
    "#创建数据集实例\n",
    "train_dataset = MyDataset(img_dir = r\"C:\\Users\\jactv\\Documents\\xwechat_files\\wxid_kq34yw840or812_f008\\msg\\file\\2025-04\\Car\\imgs\\train\", mask_dir = r\"C:\\Users\\jactv\\Documents\\xwechat_files\\wxid_kq34yw840or812_f008\\msg\\file\\2025-04\\Car\\masks\\train\", transform = transform)\n",
    "val_dataset = MyDataset(img_dir = r\"C:\\Users\\jactv\\Documents\\xwechat_files\\wxid_kq34yw840or812_f008\\msg\\file\\2025-04\\Car\\imgs\\val\", mask_dir = r\"C:\\Users\\jactv\\Documents\\xwechat_files\\wxid_kq34yw840or812_f008\\msg\\file\\2025-04\\Car\\masks\\val\", transform = transform)\n",
    "#创建数据加载器（shuffle=True 训练打乱 不打乱数据）\n",
    "train_loader = DataLoader(train_dataset, batch_size = 8, shuffle = True)\n",
    "val_loader = DataLoader(val_dataset, batch_size = 8, shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "dec28e8c-387a-4144-91e8-666e76ca0546",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "device = torch.device(\"cuda\"if torch.cuda.is_available() else \"cpu\")\n",
    "#二分类任务\n",
    "model = ResNetUNet(n_class=1).to(device)\n",
    "#二元交叉损失函数，适用于二分类任务\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "#Adam 自适应优化算法\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "\n",
    "#训练10次\n",
    "num_epochs = 10\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c18f565-a2f3-4aad-9012-57f0ef2a0e95",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63e11c43-0149-41ba-8465-5f4f966a0448",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AlexNet",
   "language": "python",
   "name": "alexnet"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
